#coding = utf-8

'''
HDenseUNET训练
'''

from network.HDenseUnet import dense_rnn_net
from dataset.IRCAD_HDENSE_UNET import IRCAD_DATA
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
import torch
from tqdm import tqdm
import numpy as np

def train(data_loader, model, criterion, optimizer):
    loss_list = []
    tbar = tqdm(data_loader, ascii=True, desc='train', dynamic_ncols=True)
    for batch_idx, (data, mask) in enumerate(tbar):
        data = data.cuda()
        mask = mask.cuda()

        print(data.shape)

        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, mask)
        loss.backward()
        optimizer.step()

        loss_list.append(loss.item())

        tbar.set_postfix(loss=f'{loss.item():.5f}')
    return loss_list





def train_dense_unet():
    ircad = IRCAD_DATA(origion_data_path="/datasets/3Dircadb/origion", hdense_data_path="/datasets/3Dircadb/HDenseNet",
                       input_size=224, input_cols=8, mean=0)

    subset = ircad.train_dataset
    model = dense_rnn_net(8).cuda()

    unet_2d_model = torch.load("/home/diaozhaoshuo/log/BeliefFunctionNN/hdenseunet/dense_unet_2d/epoll_357.pkl")
    unet_3d_model = torch.load("/home/diaozhaoshuo/log/BeliefFunctionNN/hdenseunet/dense_unet_3d/epoll_399.pkl")

    model.dense2d.load_state_dict(unet_2d_model.state_dict())
    model.dense3d.load_state_dict(unet_3d_model.dense_unet_3d.state_dict())
    model.finalConv3d1.load_state_dict(unet_3d_model.finalConv3d1.state_dict())
    model.finalBn.load_state_dict(unet_3d_model.finalBn.state_dict())
    model.finalAc.load_state_dict(unet_3d_model.finalAc.state_dict())
    model.finalConv3d2.load_state_dict(unet_3d_model.finalConv3d2.state_dict())

    data_loader = DataLoader(subset, batch_size=1, shuffle=True)
    opt = torch.optim.SGD(model.parameters(), lr=1e-4, momentum=0.9)
    loss = torch.nn.CrossEntropyLoss(weight=torch.tensor((0.78, 0.65, 8.57), device='cuda'), reduction='mean')
    model_save_path = "/home/diaozhaoshuo/log/BeliefFunctionNN/hdenseunet/fusion"

    for i in range(200):
        print("-------------------epoll{}---------------------------".format(i))

        model.train()
        torch.set_grad_enabled(True)
        loss_list = train(data_loader=data_loader, criterion=loss, optimizer=opt, model=model)
        torch.save(model, "{}/epoll_{}.pkl".format(model_save_path, str(i).zfill(3)))
        print()
        print("train loss : {}".format(np.mean(np.array(loss_list))))
        print()





if __name__ == '__main__':
    train_dense_unet()